#include "mpi.h"
#include <math.h>
#include <stdio.h>
#include <stdlib.h>

/* Externally defined global values and subroutines */
#define MAXN 1000000
extern int N;            // Number of particles (N<MAXN) is globally defined
extern double x[MAXN];   // Positions of particles are globally defined
extern double move();    // Routine for changing position of particles
                         // based on current position and applied force
int root = 0;


void update(){
double totalForce[N];
int tasksize, remainingparticles, taskstart, taskend;
int i,j;
int *rcvcounts, *displs;
double *resultdata;
    
	// Inform all the processors about the new x array from the root processor
    MPI_Bcast(x, N, MPI_DOUBLE, root, MPI_COMM_WORLD);

    // Compute the task size and the remaining particles
    remainingparticles = N % numprocs;
    tasksize = N / numprocs;
    
    // Determine the start index and end index in the x array
    taskstart = procnum * (N / numprocs);
    if( procnum == numprocs - 1){
    	tasksize = tasksize + remainingparticles;
    	taskend = taskstart + tasksize;  // More particles for the last processor
    }
    else{
    	taskend = taskstart + tasksize;
    }
    
    // Memory allocated for the use of gathering function
    displs = (int *) malloc(numprocs * sizeof(int));
    rcvcounts = (int *) malloc(numprocs * sizeof(int));
    resultdata = (double *) malloc(tasksize*sizeof(double));

    /* Determine the displs array and receive counts array */
    for (i = 0; i < numprocs; i++)
    {
        displs[i] = i * (N / numprocs);  // all the displacement are the same
        // for the last processor, receve count is more than others
        if(i == numprocs -1) {
            rcvcounts[i] = (N / numprocs) + (N % numprocs);
        } else {
        	// others process get the same task size
            rcvcounts[i] = tasksize;
        }
    }

   /* Loop over all particles to find force on each particle */
    for(i = taskstart; i < taskend; i++)
    {
        totalForce[i] = 0.0;
        // for particle i, sum forces from all other particles
        for(j = 0; j < N; j++)
        {
            if (j < i)
                totalForce[i] -= 1.0 / pow(x[i] - x[j],2);
            if (j > i)
                totalForce[i] += 1.0 / pow(x[i] - x[j],2);
        }
    }

    /* Loop over all particles to update position of each particle */
    for(i = 0; i < taskend - taskstart; i++)
    {
        /* Update position of particles to the resultdata array */
        resultdata[i] = move(x[taskstart + i],totalForce[taskstart + i]);
    }

    /* use MPI_Gatherv so that the root process can receive  different amount of data from each process */
    MPI_Gatherv(resultdata, tasksize, MPI_DOUBLE, x, rcvcounts, displs, MPI_DOUBLE, root, MPI_COMM_WORLD);

}

/******************************************************************************/
/* Update positions of particles using force from nearest neighbour particles */
/******************************************************************************/

void updateNN() {
double totalForce[N];
int tasksize, remainingparticles, taskstart, taskend;
int i,j;
int *rcvcounts, *displs;
double *resultdata;
    
	// Inform all the processors about the new x array from the root processor
    MPI_Bcast(x, N, MPI_DOUBLE, root, MPI_COMM_WORLD);

    // Compute the task size and the remaining particles
    remainingparticles = N % numprocs;
    tasksize = N / numprocs;
    
    // Determine the start index and end index in the x array
    taskstart = procnum * (N / numprocs);
    if( procnum == numprocs - 1){
    	tasksize = tasksize + remainingparticles;
    	taskend = taskstart + tasksize;  // More particles for the last processor
    }
    else{
    	taskend = taskstart + tasksize;
    }
    
    // Memory allocated for the use of gathering function
    displs = (int *) malloc(numprocs * sizeof(int));
    rcvcounts = (int *) malloc(numprocs * sizeof(int));
    resultdata = (double *) malloc(tasksize*sizeof(double));

    /* Determine the displs array and receive counts array */
    for (i = 0; i < numprocs; i++)
    {
        displs[i] = i * (N / numprocs);  // all the displacement are the same
        // for the last processor, receve count is more than others
        if(i == numprocs -1) {
            rcvcounts[i] = (N / numprocs) + (N % numprocs);
        } else {
        	// others process get the same task size
            rcvcounts[i] = tasksize;
        }
    }

    /* Loop over all particles to find force on each particle */
    for(i = taskstart; i < taskend; i++)
    {
        totalForce[i] = 0.0;
        /* Sum forces over all particles to get force on particle i */
        if(i > 0)
            totalForce[i] -= 1.0 / pow(x[i] - x[i-1],2);
        if (i < N-1)
            totalForce[i] += 1.0 / pow(x[i] - x[i+1],2);
    }

    /* Loop over all particles to update position of each particle */
    for(i = 0; i < taskend - taskstart; i++)
    {
        /* Update position of particles to the resultdata array */
        resultdata[i] = move(x[taskstart + i],totalForce[taskstart + i]);
    }

//    MPI_Barrier(MPI_COMM_WORLD);

    /* use MPI_Gatherv to gather result so that the root process can receive  varying count of data from each process */
    MPI_Gatherv(resultdata, tasksize, MPI_DOUBLE, x, rcvcounts, displs, MPI_DOUBLE, root, MPI_COMM_WORLD);
    
}

/***********************************************************/
/* Find average distance between neighbouring particles.   */
/* NOTE: In 1-dimension there is an easier way to do this, */
/* but assume that we have to use the following code.      */
/***********************************************************/
double averageDistance() {
    int i;
    double distance;

    /* Loop over all particles */
    distance = 0.0;
    for (i=1;i<N;i++) {
        distance += x[i] - x[i-1];
    }
    /* Compute average distance */
    distance = distance / (N-1);
    return distance;
}
